import pandas as pd
import requests
import logging
import os
import time
from tqdm import tqdm
from zhipuai import ZhipuAI

'通过智谱API调用glm-4-plus来查询教练名称'

def setup_logging():
    """设置日志记录"""
    log_dir = 'logs'
    if not os.path.exists(log_dir):
        os.makedirs(log_dir)

    log_file = f'{log_dir}/coaches_info_generation.log'
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(log_file, encoding='utf-8'),
            logging.StreamHandler()
        ]
    )


def call_zhipu_api(country, year, sport, coaches_cache):
    """使用智谱AI SDK调用API获取教练信息"""
    key = f"{country}_{year}_{sport}"
    if key in coaches_cache:
        return coaches_cache[key]

    client = ZhipuAI(api_key="77c93f2f91e14dd885888de3f3b5f26a.PKBkd1xAPA3dTzkp")  # 替换为你的API密钥

    messages = [
        {
            "role": "system",
            "content": "你是一个体育教练信息查询助手。"
        },
        {
            "role": "user",
            "content": f"请告诉我 {country} 在 {year} 年的 {sport} 项目的主教练是谁？"
        }
    ]

    try:
        response = client.chat.completions.create(
            model="glm-4-plus",
            messages=messages,
            temperature=0.1,
            top_p=0.1,
        )

        coach_name = response.choices[0].message.content.strip()
        logging.info(f"获取到的教练名: {coach_name}")

        # 缓存教练信息
        coaches_cache[key] = coach_name
        return coach_name

    except Exception as e:
        logging.error(f"API调用错误: {str(e)}")
        return '错误'


def load_data(filepath):
    """读取CSV文件"""
    if os.path.exists(filepath):
        logging.info(f"加载数据: {filepath}")
        return pd.read_csv(filepath)
    else:
        logging.error(f"文件不存在: {filepath}")
        return None


def save_results(results, output_filepath):
    """保存结果到CSV"""
    logging.info(f"保存结果到: {output_filepath}")
    results.to_csv(output_filepath, index=False)


def main():
    setup_logging()
    input_filepath = 'summerOly_athletes.csv'
    output_filepath = 'coaches_info.csv'
    start_index = 0

    # 检查是否有已保存的结果
    if os.path.exists(output_filepath):
        existing_results = pd.read_csv(output_filepath)
        start_index = len(existing_results)
        logging.info(f"从索引 {start_index} 继续处理")

    data = load_data(input_filepath)
    if data is None:
        return

    results = []
    coaches_cache = {}

    # 使用 tqdm 显示进度条
    for index in tqdm(range(start_index, len(data)), desc="Processing", unit="record"):
        row = data.iloc[index]
        country = row['Team']
        year = row['Year']
        sport = row['Sport']

        logging.info(f"处理 {country}, {year}, {sport}...")
        coach_name = call_zhipu_api(country, year, sport, coaches_cache)

        results.append({
            "Country": country,
            "Year": year,
            "Sport": sport,
            "Coach": coach_name
        })

        # 每处理100条记录保存一次
        if (index + 1) % 100 == 0:
            logging.info(f"已处理 {index + 1} 条记录，保存中...")
            save_results(pd.DataFrame(results), output_filepath)

        # 控制请求频率，避免API限制
        time.sleep(1)

    # 保存所有结果
    logging.info("最终保存结果...")
    save_results(pd.DataFrame(results), output_filepath)
    logging.info("处理完成。")


if __name__ == "__main__":
    main()
